package com.lw.leetcode.tree.b;

/**
 * Created with IntelliJ IDEA.
 * 208. 实现 Trie (前缀树)
 * 剑指 Offer II 062. 实现前缀树
 *
 * @author liw
 * @version 1.0
 * @date 2021/11/3 17:46
 */
public class Trie {
    private Tree[] roots;

    public Trie() {
        roots = new Tree[26];
    }

    public void insert(String word) {
        char[] chars = word.toCharArray();
        insert(roots, chars, 0);
    }

    public boolean search(String word) {
        char[] chars = word.toCharArray();
        return find(roots, chars, 0);
    }

    public boolean startsWith(String prefix) {
        char[] chars = prefix.toCharArray();
        return findStartsWith(roots, chars, 0);
    }

    private void insert(Tree[] nodes, char[] chars, int index) {
        char aChar = chars[index];
        Tree node = nodes[aChar - 'a'];
        if (node == null) {
            node = new Tree();
            nodes[aChar - 'a'] = node;
        }
        node.have = true;
        if (index == chars.length - 1) {
            node.end = true;
            return;
        }
        insert(node.nexts, chars, index + 1);
    }

    private boolean find(Tree[] nodes, char[] chars, int index) {
        char c = chars[index];
        Tree node = nodes[c - 'a'];
        if (node != null && node.have) {
            if (node.end && chars.length - 1 == index) {
                return true;
            }
            if (chars.length - 1 == index) {
                return false;
            }
            return find(node.nexts, chars, index + 1);
        }
        return false;
    }

    private boolean findStartsWith(Tree[] nodes, char[] chars, int index) {
        char c = chars[index];
        Tree node = nodes[c - 'a'];
        if (node != null && node.have) {
            if (chars.length - 1 == index) {
                return true;
            }
            return findStartsWith(node.nexts, chars, index + 1);
        }
        return false;
    }

    private class Tree {
        private boolean have;
        private boolean end;
        private Tree[] nexts = new Tree[26];
    }
}
